# 导入torch工具包
import torch

# 选定预训练模型的来源
source = "huggingface/pytorch-transformers"

# 选定加载模型的哪一部分, 这里我么选定模型的映射器
part = "tokenizer"

# 选定加载的预训练模型的名称
model_name = "bert-base-chinese"

# 直接通过torch.hub加载预训练模型的映射器
tokenizer = torch.hub.load(source, part, model_name)

# 加载不带头的预训练模型
# part = "model"
# model = torch.hub.load(source, part, model_name)

# 加载带有语言模型头的预训练模型
# part = "modelWithLMHead"
# lm_model = torch.hub.load(source, part, model_name)

# 加载带有类模型头的预训练模型
# part = "modelForSequenceClassification"
# classification_model = torch.hub.load(source, part, model_name)

# 加载带有问答模型头的预训练模型
part = "modelForQuestionAnswering"
qa_model = torch.hub.load(source, part, model_name)

# 输入一行中文文本
input_text = "人生该如何起头"

# 第一步使用映射器进行数值的映射
indexed_tokens = tokenizer.encode(input_text)

# 打印一下映射后的编码结构
# print("indexed_tokens:", indexed_tokens)

# 将映射后的结构利用torch.tensor()进行封装, 才可以传入模型
tokens_tensor = torch.tensor([indexed_tokens])

# 使用不带头的预训练模型model来获取结果
# with torch.no_grad():
#     encoded_layers, _ = model(tokens_tensor)

# print("不带头的模型输出张量结果:", encoded_layers)
# print("不带头的模型输出张量的尺寸:", encoded_layers.shape)

# 使用带有语言模型头的预训练模型获得结果
# with torch.no_grad():
#     lm_output = lm_model(tokens_tensor)

# print("带语言模型头的模型输出张量的结果:", lm_output)
# print("带语言模型头的模型输出张量的形状:", lm_output[0].shape)

# 使用带有分类模型头的预训练模型获得结果
# with torch.no_grad():
#     classification_output = classification_model(tokens_tensor)

# print("带分类模型头的模型输出张量结果:", classification_output)
# print("带分类模型头的模型输出张量尺寸:", classification_output[0].shape)

# 初始化两个句子对
input_text1 = "我家的小狗是黑色的"
input_text2 = "我家的小狗是什么颜色的呢?"

# 利用映射器得到句子对的数字化张量
indexed_tokens = tokenizer.encode(input_text1, input_text2)
print("句子对的indexed_tokens:", indexed_tokens)

# 利用0来标志第一条语句的位置, 利用1来标志第二条语句的位置
segments_ids = [0]*11 + [1]*14

# 将位置编码和句子编码进行torch.tensor的封装
segments_tensor = torch.tensor([segments_ids])
tokens_tensor = torch.tensor([indexed_tokens])

# 利用带头问答模型头的预训练模型得到结果
with torch.no_grad():
    start_logits, end_logits = qa_model(tokens_tensor, token_type_ids=segments_tensor)

print("带问答模型头的模型输出张量:", (start_logits, end_logits))
print("带问答模型头的模型输出张量的尺寸:", (start_logits.shape, end_logits.shape))































